//
//  SpMV.metal
//  Test
//
//  Created by 张木林 on 2/18/25.
//

#include <metal_stdlib>
#include <metal_simdgroup>
using namespace metal;



#define BLOCK_SIZE 32











//// ========================================
//// CSR Kernel Function
//// ========================================
//kernel void spmv_csr(
//    device const float* __restrict values [[buffer(0)]],
//    device const int* __restrict columnIndices [[buffer(1)]],
//    device const int* __restrict rowPointers [[buffer(2)]],
//    device const float* __restrict x [[buffer(3)]],
//    device float* __restrict y [[buffer(4)]],
//    uint gid [[thread_position_in_grid]]
//) {
//    if (gid >= rowPointers[0]) return;
//
//    float sum = 0.0f;
//    int row_start = rowPointers[gid];
//    int row_end = rowPointers[gid + 1];
//
//    for (int i = row_start; i < row_end; i++) {
//        sum += values[i] * x[columnIndices[i]];
//    }
//    y[gid] = sum;
//}
//
//// ========================================
//// COO Kernel Function
//// ========================================
//kernel void spmv_coo(
//    device const float* __restrict values [[buffer(0)]],
//    device const int* __restrict rowIndices [[buffer(1)]],
//    device const int* __restrict columnIndices [[buffer(2)]],
//    device const float* __restrict x [[buffer(3)]],
//    device float* __restrict y [[buffer(4)]],
//    uint gid [[thread_position_in_grid]]
//) {
//    if (gid >= rowIndices[0]) return;
//
//    int row = rowIndices[gid];
//    int col = columnIndices[gid];
//    atomic_fetch_add_explicit((device atomic_float*)&y[row], values[gid] * x[col], memory_order_relaxed);
//}
//
//// ========================================
//// CSC Kernel Function
//// ========================================
//kernel void spmv_csc(
//    device const float* __restrict values [[buffer(0)]],
//    device const int* __restrict rowIndices [[buffer(1)]],
//    device const int* __restrict columnPointers [[buffer(2)]],
//    device const float* __restrict x [[buffer(3)]],
//    device float* __restrict y [[buffer(4)]],
//    uint gid [[thread_position_in_grid]]
//) {
//    if (gid >= columnPointers[0]) return;
//
//    int col_start = columnPointers[gid];
//    int col_end = columnPointers[gid + 1];
//
//    for (int i = col_start; i < col_end; i++) {
//        int row = rowIndices[i];
//        atomic_fetch_add_explicit((device atomic_float*)&y[row], values[i] * x[gid], memory_order_relaxed);
//    }
//}
//
//// ========================================
//// HYB Kernel Function
//// ========================================
//kernel void spmv_hyb(
//    device const float* __restrict csr_values [[buffer(0)]],
//    device const int* __restrict csr_columnIndices [[buffer(1)]],
//    device const int* __restrict csr_rowPointers [[buffer(2)]],
//    device const float* __restrict ell_values [[buffer(3)]],
//    device const int* __restrict ell_columnIndices [[buffer(4)]],
//    device const float* __restrict x [[buffer(5)]],
//    device float* __restrict y [[buffer(6)]],
//    uint gid [[thread_position_in_grid]]
//) {
//    if (gid >= csr_rowPointers[0]) return;
//
//    // Process CSR part
//    float sum = 0.0f;
//    int row_start = csr_rowPointers[gid];
//    int row_end = csr_rowPointers[gid + 1];
//
//    for (int i = row_start; i < row_end; i++) {
//        sum += csr_values[i] * x[csr_columnIndices[i]];
//    }
//
//    // Process ELL part
//    int ellWidth = 32; // 设定固定的 ELL 列宽
//    for (int i = 0; i < ellWidth; i++) {
//        int col = ell_columnIndices[gid * ellWidth + i];
//        if (col >= 0) {
//            sum += ell_values[gid * ellWidth + i] * x[col];
//        }
//    }
//
//    y[gid] = sum;
//}



// ✅ **修正后的 CSR Kernel**
kernel void spmv_csr(device const float *values [[buffer(0)]],
                     device const int *columnIndices [[buffer(1)]],
                     device const int *rowPointers [[buffer(2)]],
                     device const float *x [[buffer(3)]],
                     device float *y [[buffer(4)]],
                     uint row [[thread_position_in_grid]],
                     uint local_id [[thread_position_in_threadgroup]]) { // ✅ 这里定义局部线程 ID
    float sum = 0.0;
    
    // ✅ 使用 `threadgroup_memory` 缓存 `x`
    threadgroup float shared_x[128];

    for (int i = rowPointers[row]; i < rowPointers[row + 1]; i++) {
        int colIdx = columnIndices[i];

        // ✅ 修正 `thread_position_in_threadgroup`
        if (local_id < 128) {
            shared_x[local_id] = x[colIdx];
        }

        // ✅ **最终修正 `threadgroup_barrier()` 语法**
        threadgroup_barrier(mem_flags::mem_threadgroup);
        
        sum += values[i] * shared_x[colIdx % 128];  // 访问共享内存
    }

    y[row] = sum;
}

// ✅ **优化的 COO Kernel**
kernel void spmv_coo(device const float *values [[buffer(0)]],
                     device const int *rowIndices [[buffer(1)]],
                     device const int *colIndices [[buffer(2)]],
                     device const float *x [[buffer(3)]],
                     device float *y [[buffer(4)]],
                     uint i [[thread_position_in_grid]]) {
    
    float val = values[i];
    int row = rowIndices[i];
    int col = colIndices[i];
    
    // ✅ 避免多个线程并发写入 `y`，使用 `atomic_add`
    atomic_fetch_add_explicit((device atomic_float*)&y[row], val * x[col], memory_order_relaxed);
}

// ✅ **优化的 CSC Kernel**
kernel void spmv_csc(device const float *values [[buffer(0)]],
                     device const int *rowIndices [[buffer(1)]],
                     device const int *colPointers [[buffer(2)]],
                     device const float *x [[buffer(3)]],
                     device float *y [[buffer(4)]],
                     uint col [[thread_position_in_grid]]) {
    
    float xVal = x[col];
    
    for (int i = colPointers[col]; i < colPointers[col + 1]; i++) {
        int row = rowIndices[i];
        
        // ✅ 避免多个线程并发写入 `y`，使用 `atomic_add`
        atomic_fetch_add_explicit((device atomic_float*)&y[row], values[i] * xVal, memory_order_relaxed);
    }
}

// ✅ **优化的 HYB Kernel**
kernel void spmv_hyb(device const float *csrValues [[buffer(0)]],
                     device const int *csrColumnIndices [[buffer(1)]],
                     device const int *csrRowPointers [[buffer(2)]],
                     device const float *ellValues [[buffer(3)]],
                     device const int *ellColumnIndices [[buffer(4)]],
                     device const float *x [[buffer(5)]],
                     device float *y [[buffer(6)]],
                     uint row [[thread_position_in_grid]]) {
    
    float sum = 0.0;
    
    // ✅ **CSR 部分计算**
    for (int i = csrRowPointers[row]; i < csrRowPointers[row + 1]; i++) {
        sum += csrValues[i] * x[csrColumnIndices[i]];
    }
    
    // ✅ **ELL 部分计算**
    for (int i = 0; i < 128; i++) {  // 限制 `ellWidth` 避免过大
        int colIdx = ellColumnIndices[row * 128 + i];
        if (colIdx == -1) break;
        sum += ellValues[row * 128 + i] * x[colIdx];
    }
    
    y[row] = sum;
}
